"""Micro-benchmark of PR #947's per-query prepared-metadata bookkeeping.

Run with: python3 bench.py    (no dependencies, no driver install needed)

_class_attribute, _prepared_metadata_cache_config,
_prepared_metadata_protocol_handler_snapshot and _prepared_metadata_decoder_context
are copied VERBATIM from the PR (cassandra/protocol.py, cassandra/cluster.py).
Only the surrounding driver classes are stubbed, shaped to match the real ones:
ResultMessage's full class-attribute set, FastResultMessage's code_to_type,
_ProtocolHandler / CythonProtocolHandler, 24 CQL type codes, 10 opcodes.

This isolates the cost of the new code paths only. It is NOT an end-to-end driver
benchmark -- real per-query cost includes encoding, I/O and row decoding on top.
"""
import timeit
from threading import Lock
from types import MappingProxyType

# ---------------------------------------------------------------- stub driver

class _MessageType(object):
    tracing = False
    custom_payload = None
    warnings = None


# 24 CQL type codes, as in cassandra/type_codes.py
_TYPE_CODES = {i: type('T%d' % i, (), {}) for i in range(24)}


class ResultMessage(_MessageType):
    """Mirrors cassandra.protocol.ResultMessage's class-attribute surface."""
    opcode = 0x08
    name = 'RESULT'
    kind = None
    results = None
    paging_state = None
    type_codes = _cqltypes_by_code = _TYPE_CODES
    _FLAGS_GLOBAL_TABLES_SPEC = 0x0001
    _HAS_MORE_PAGES_FLAG = 0x0002
    _NO_METADATA_FLAG = 0x0004
    _CONTINUOUS_PAGING_FLAG = 0x40000000
    _CONTINUOUS_PAGING_LAST_FLAG = 0x80000000
    _METADATA_ID_FLAG = 0x0008
    column_names = None
    column_types = None
    parsed_rows = None
    continuous_paging_seq = None
    continuous_paging_last = None
    new_keyspace = None
    column_metadata = None
    query_id = None
    bind_metadata = None
    pk_indexes = None
    schema_change_event = None
    is_lwt = False

    def __init__(self, kind):
        self.kind = kind

    def recv(self, *a): pass
    @classmethod
    def recv_body(cls, *a): pass
    def recv_results_rows(self, *a): pass
    def recv_results_prepared(self, *a): pass
    def recv_results_metadata(self, *a): pass
    def recv_prepared_metadata(self, *a): pass
    def recv_results_schema_change(self, *a): pass
    @classmethod
    def read_type(cls, *a): pass
    @staticmethod
    def recv_row(*a): pass


class FastResultMessage(ResultMessage):
    code_to_type = dict((v, k) for k, v in ResultMessage.type_codes.items())
    def recv_results_rows(self, *a): pass


# 10 opcodes, as in _ProtocolHandler.message_types_by_opcode
_OPCODES = {i: type('M%d' % i, (), {}) for i in range(10)}


# ----------------------------- verbatim from PR: cassandra/protocol.py

_PREPARED_METADATA_RESULT_ATTRIBUTES = (
    'opcode', 'name', 'kind', 'results', 'paging_state',
    '_FLAGS_GLOBAL_TABLES_SPEC', '_HAS_MORE_PAGES_FLAG', '_NO_METADATA_FLAG',
    '_CONTINUOUS_PAGING_FLAG', '_CONTINUOUS_PAGING_LAST_FLAG',
    '_METADATA_ID_FLAG', 'column_names', 'column_types', 'parsed_rows',
    'continuous_paging_seq', 'continuous_paging_last', 'new_keyspace',
    'column_metadata', 'query_id', 'bind_metadata', 'pk_indexes',
    'schema_change_event', 'is_lwt', 'tracing', 'custom_payload', 'warnings',
    '__init__', 'recv', 'recv_body', 'recv_results_rows',
    'recv_results_prepared', 'recv_results_metadata', 'recv_prepared_metadata',
    'recv_results_schema_change', 'read_type', 'recv_row',
)


def _class_attribute(cls, name):
    for base in cls.__mro__:
        if name in vars(base):
            return vars(base)[name]
    return None


_PREPARED_METADATA_CONFIG_UNSET = object()


def _prepared_metadata_cache_config(
        handler, result_message_type,
        message_types=_PREPARED_METADATA_CONFIG_UNSET,
        type_codes=_PREPARED_METADATA_CONFIG_UNSET,
        code_to_type=_PREPARED_METADATA_CONFIG_UNSET,
        decode_message=_PREPARED_METADATA_CONFIG_UNSET,
        result_attributes=_PREPARED_METADATA_CONFIG_UNSET,
        column_encryption_policy=_PREPARED_METADATA_CONFIG_UNSET):
    message_types = (
        handler.message_types_by_opcode
        if message_types is _PREPARED_METADATA_CONFIG_UNSET else message_types)
    type_codes = (
        result_message_type.type_codes
        if type_codes is _PREPARED_METADATA_CONFIG_UNSET else type_codes)
    code_to_type = (
        getattr(result_message_type, 'code_to_type', {})
        if code_to_type is _PREPARED_METADATA_CONFIG_UNSET else code_to_type)
    decode_message = (
        _class_attribute(handler, 'decode_message')
        if decode_message is _PREPARED_METADATA_CONFIG_UNSET else decode_message)
    result_attributes = (
        tuple((name, _class_attribute(result_message_type, name))
              for name in _PREPARED_METADATA_RESULT_ATTRIBUTES)
        if result_attributes is _PREPARED_METADATA_CONFIG_UNSET
        else result_attributes)
    column_encryption_policy = (
        handler.column_encryption_policy
        if column_encryption_policy is _PREPARED_METADATA_CONFIG_UNSET
        else column_encryption_policy)
    return (tuple(message_types.items()), result_message_type,
            tuple(type_codes.items()), tuple(code_to_type.items()),
            decode_message, result_attributes, column_encryption_policy)


class _ProtocolHandler(object):
    message_types_by_opcode = _OPCODES
    column_encryption_policy = None
    _prepared_metadata_cache_token = object()

    @classmethod
    def decode_message(cls, *a): pass


_ProtocolHandler._prepared_metadata_cache_config = \
    _prepared_metadata_cache_config(_ProtocolHandler, ResultMessage)


def cython_protocol_handler():
    class CythonProtocolHandler(_ProtocolHandler):
        my_opcodes = _ProtocolHandler.message_types_by_opcode.copy()
        my_opcodes[FastResultMessage.opcode] = FastResultMessage
        message_types_by_opcode = my_opcodes
        _prepared_metadata_cache_token = object()
    CythonProtocolHandler._prepared_metadata_cache_config = \
        _prepared_metadata_cache_config(CythonProtocolHandler, FastResultMessage)
    return CythonProtocolHandler


ProtocolHandler = cython_protocol_handler()

# ------------------------------ verbatim from PR: cassandra/cluster.py

_PREPARED_METADATA_HANDLER_SNAPSHOTS = {}
_PREPARED_METADATA_HANDLER_SNAPSHOT_LOCK = Lock()


def _callable_identity(value):
    return getattr(value, '__func__', value)


def _prepared_metadata_protocol_handler_snapshot(protocol_handler):
    try:
        handler_vars = vars(protocol_handler)
    except TypeError:
        return protocol_handler

    if handler_vars.get('_prepared_metadata_handler_snapshot'):
        return protocol_handler

    handler_token = handler_vars.get('_prepared_metadata_cache_token')
    if handler_token is None:
        return protocol_handler

    try:
        message_types = dict(protocol_handler.message_types_by_opcode)
        result_message_type = message_types[ResultMessage.opcode]
        if not isinstance(result_message_type, type):
            return protocol_handler
        type_codes = dict(getattr(result_message_type, 'type_codes', {}))
        code_to_type = dict(getattr(result_message_type, 'code_to_type', {}))
        decode_message = _class_attribute(protocol_handler, 'decode_message')
        result_attributes = tuple(
            (name, _class_attribute(result_message_type, name))
            for name in _PREPARED_METADATA_RESULT_ATTRIBUTES
        )
        column_encryption_policy = getattr(
            protocol_handler, 'column_encryption_policy', None)
        handler_config = _prepared_metadata_cache_config(
            protocol_handler, result_message_type, message_types, type_codes,
            code_to_type, decode_message, result_attributes,
            column_encryption_policy,
        )
        canonical_config = handler_vars.get('_prepared_metadata_cache_config')
        if (canonical_config is None
                or handler_config[:-1] != canonical_config[:-1]
                or column_encryption_policy is not canonical_config[-1]):
            return protocol_handler
        cache_key = (protocol_handler, handler_token, handler_config[:-1],
                     id(column_encryption_policy))
        hash(cache_key)
    except Exception:
        return protocol_handler

    with _PREPARED_METADATA_HANDLER_SNAPSHOT_LOCK:
        try:
            cached = _PREPARED_METADATA_HANDLER_SNAPSHOTS.get(protocol_handler)
            if cached is not None and cached[0] == cache_key:
                return cached[1]
        except Exception:
            return protocol_handler

        try:
            result_attrs = dict(result_attributes)
            result_attrs.update({
                '__module__': result_message_type.__module__,
                'type_codes': MappingProxyType(type_codes),
                '_cqltypes_by_code': MappingProxyType(type_codes),
                'code_to_type': MappingProxyType(code_to_type),
            })
            result_snapshot = type(
                '_%sPreparedMetadataSnapshot' % result_message_type.__name__,
                (result_message_type,), result_attrs)

            message_types[ResultMessage.opcode] = result_snapshot
            handler_attrs = {
                '__module__': protocol_handler.__module__,
                'message_types_by_opcode': MappingProxyType(message_types),
                'column_encryption_policy': column_encryption_policy,
                '_prepared_metadata_cache_token': object(),
                '_prepared_metadata_handler_snapshot': True,
            }
            if decode_message is not None:
                handler_attrs['decode_message'] = decode_message
            snapshot = type(
                '%sPreparedMetadataSnapshot' % protocol_handler.__name__,
                (protocol_handler,), handler_attrs)
        except Exception:
            return protocol_handler
        _PREPARED_METADATA_HANDLER_SNAPSHOTS[protocol_handler] = (
            cache_key, snapshot)
        return snapshot


def _prepared_metadata_decoder_context(protocol_handler, cluster_context):
    if cluster_context is None:
        return None

    try:
        handler_vars = vars(protocol_handler)
    except TypeError:
        return None
    if not handler_vars.get('_prepared_metadata_handler_snapshot'):
        return None

    result_message_type = \
        protocol_handler.message_types_by_opcode.get(ResultMessage.opcode)
    type_codes = getattr(result_message_type, 'type_codes', {})
    code_to_type = getattr(result_message_type, 'code_to_type', {})
    column_encryption_policy = getattr(
        protocol_handler, 'column_encryption_policy', None)
    return (
        handler_vars['_prepared_metadata_cache_token'],
        protocol_handler,
        cluster_context,
        _callable_identity(protocol_handler.decode_message),
        result_message_type,
        _callable_identity(
            getattr(result_message_type, 'recv_results_rows', None)),
        _callable_identity(getattr(result_message_type, 'read_type', None)),
        tuple(type_codes.items()),
        tuple(code_to_type.items()),
        column_encryption_policy,
        id(column_encryption_policy),
    )


# ------------------------------------------------------------------ benchmark

CLUSTER_CONTEXT = object()
# Warm the shared snapshot cache and produce a cached provenance token, i.e. the
# steady state of a Session that has already executed a prepared statement.
_SNAPSHOT = _prepared_metadata_protocol_handler_snapshot(ProtocolHandler)
CACHED_CTX = _prepared_metadata_decoder_context(_SNAPSHOT, CLUSTER_CONTEXT)


def per_query():
    """What one already-warm BoundStatement execute now costs.

    Session._create_response_future passes the SOURCE handler class, so the
    `_prepared_metadata_handler_snapshot` early return never fires and the full
    revalidation runs again. ResponseFuture.__init__ then rebuilds the decoder
    context a second time.
    """
    # Session._create_response_future
    handler = _prepared_metadata_protocol_handler_snapshot(ProtocolHandler)
    ctx = _prepared_metadata_decoder_context(handler, CLUSTER_CONTEXT)
    eligible = ctx is not None and CACHED_CTX == ctx
    # ResponseFuture.__init__ (snapshot early-returns here; context is rebuilt)
    handler2 = _prepared_metadata_protocol_handler_snapshot(handler)
    ctx2 = _prepared_metadata_decoder_context(handler2, CLUSTER_CONTEXT)
    return eligible and ctx == ctx2


assert per_query(), "steady state should be skip_meta-eligible"

N = 20000
print("per prepared-statement execute: %5.1f us"
      % (timeit.timeit(per_query, number=N) / N * 1e6))
for label, fn in (
    ("  _prepared_metadata_protocol_handler_snapshot()",
     lambda: _prepared_metadata_protocol_handler_snapshot(ProtocolHandler)),
    ("  _prepared_metadata_decoder_context()",
     lambda: _prepared_metadata_decoder_context(_SNAPSHOT, CLUSTER_CONTEXT)),
):
    print("%s: %5.1f us" % (label, timeit.timeit(fn, number=N) / N * 1e6))